# -*- coding:utf-8 -*-

import random
import optparse
import io
def extract_id(trips):
    assert type(trips) is list
    idset = []
    for trip in trips:
        assert trip['id']
        idset.append(trip['id'])
    return idset

def generate_white_black(idset, proportion=0.5):
    assert 0 < proportion < 1.0
    # generate white list
    # shuffle
    random.shuffle(idset)
    k = int(len(idset) * proportion)
    sample = random.sample(idset, k)
    black_dict = {}.fromkeys(sample)
    white_dict = {}
    for person in idset:
        if person in black_dict:
            continue
        white_dict.setdefault(person)
    return white_dict.keys(), black_dict.keys()

def swap_white_black(white, black, white_proportion=0.2, black_proportion=0.2):
    assert type(white) is list
    assert type(black) is list

    white_white, white_black = generate_white_black(white, proportion=white_proportion)
    black_white, black_black = generate_white_black(black, proportion=black_proportion)

    white_white.extend(black_white)
    white_black.extend(black_black)
    return white_white, white_black

def adjust_depart(trips, mu=0, sigma=300):
    assert type(trips) is list
    for i, trip in enumerate(trips):
        variation = int(random.gauss(mu=mu, sigma=sigma))
        assert trip['depart']
        trips[i]['depart'] = int(float(trip['depart'])) + variation
    return sorted(trips, key=lambda d: d['depart'], reverse=False)

def assign_vehicle_stop(edges, trip, probability):
    assert type(trip) is dict
    proba = random.random()
    if proba <= probability:
        # accept
        route = trip['route']
        e_id = random.randint(1, len(route) - 2)
        e_id = route[e_id]
        while edges[e_id] < 10:
            e_id = random.randint(2, len(route)) - 1
            e_id = route[e_id]
        end_pos = random.random() * (edges[e_id] - 10) + 10
        duration = random.randint(1, 30)
        #add stop attributes to trip
        # trip.setdefault("stop", {"lane": "{0}_0".format(e_id), "endPos":end_pos, "duration": duration})
        trip["stop"] = {"lane": "{0}_0".format(e_id), "endPos":end_pos, "duration": duration}
    return trip

def assign_emergency_stop(trips, black, edges, probability=0.3):
    assert type(trips) is list
    assert type(black) is list
    black_dict = {}.fromkeys(black)
    for i, trip in enumerate(trips):
        if trip['id'] in black_dict:
            trips[i] = assign_vehicle_stop(edges, trips[i], probability)
    return trips

def get_options(args=None):
    parser = optparse.OptionParser()
    parser.add_option("-n", "--network", dest="network", help="network")
    parser.add_option("-t", "--trips", dest="trips", help="origin routes which contains the depart, vclass, etc")
    parser.add_option("-o", "--output", dest="output", help="output file")
    parser.add_option("-s", "--size", dest="size", type="float", help="split size of bad driver group, valid range (0, 1)")
    parser.add_option("-p", "--proba", dest="proba", type="float", help="probability for bad driver group having a emergency stop default is 0.3")
    parser.add_option("-g", "--group-file", dest="group", help="vehicle group file good, bad")
    parser.add_option("--save", dest="save", help="save the two groups into file save, when size specified, save should be set")
    parser.add_option("--swap-white", dest="swap_white", type='float', help="the proportion of white set for swapping the default proportion is 0.0")
    parser.add_option("--swap-black", dest="swap_black", type='float', help="the proportion of white set for swapping the default proportion is 0.0")
    (options, args) = parser.parse_args(args=args)

    if not options.network and not options.trips and not options.output and not options.save:
        parser.print_help()
        exit(1)

    if options.size <= 0 or options.size >= 1 and options.group:
        parser.print_help()
        exit(1)
    if 0 < options.size < 1 and not options.save:
        parser.print_help()
        exit(1)
    if not options.proba:
        options.proba = 0.3
    if options.group:
        if not options.swap_white:
            options.swap_white = 0.0
        if not options.swap_black:
            options.swap_black = 0.0
        if options.swap_white > 0 and options.swap_black > 0:
            if not options.save:
                parser.print_help()
                exit(1)

    return options

def main(options=None):
    edges = io.load_net(options.network)
    trips = io.load_routes(options.trips)
    white_group = None
    black_group = None
    idset = extract_id(trips)
    if  0 < options.size < 1:
        white_group, black_group = generate_white_black(idset, proportion=options.size)
        io.save_group(options.save, white_group, black_group)
    else:
        white_group, black_group = io.load_group(options.group)
    if 0 < options.swap_white < 1 and 0 < options.swap_black < 1:
        white_group, black_group = swap_white_black(white_group, black_group, white_proportion=options.swap_white, black_proportion=options.swap_black)
        io.save_group(options.save, white_group, black_group)

    trips = adjust_depart(trips)
    trips = assign_emergency_stop(trips, black_group, edges, probability=options.proba)
    io.write_routes_text(options.output, trips)


if __name__ == '__main__':
    main(get_options())